#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging
from aliyunsdkcore.client import AcsClient
from aliyunsdkcore.acs_exception.exceptions import ClientException, ServerException
from aliyunsdkecs.request.v20140526.DescribeSecurityGroupAttributeRequest import DescribeSecurityGroupAttributeRequest
from aliyunsdkecs.request.v20140526.RevokeSecurityGroupRequest import RevokeSecurityGroupRequest
from aliyunsdkecs.request.v20140526.AuthorizeSecurityGroupRequest import AuthorizeSecurityGroupRequest

logger = logging.getLogger(__name__)

class AlibabaCloudSecurityGroupManager:
    def __init__(self, access_key_id: str, access_key_secret: str, region_id: str):
        """
        初始化阿里云安全组管理器
        
        Args:
            access_key_id: 阿里云访问密钥ID
            access_key_secret: 阿里云访问密钥
            region_id: 区域ID
        """
        self.client = AcsClient(access_key_id, access_key_secret, region_id)
        self.region_id = region_id
        
    def get_existing_rules(self, security_group_id: str, description: str) -> list:
        """
        获取指定安全组中由本程序添加的规则
        
        Args:
            security_group_id: 安全组ID
            description: 规则描述标识
            
        Returns:
            匹配的规则列表
        """
        try:
            request = DescribeSecurityGroupAttributeRequest()
            request.set_SecurityGroupId(security_group_id)
            request.set_accept_format('json')
            
            response = self.client.do_action_with_exception(request)
            # 解析响应，查找匹配描述的规则
            # 这里需要根据实际API响应结构调整
            import json
            result = json.loads(response)
            
            matched_rules = []
            if 'Permissions' in result and 'Permission' in result['Permissions']:
                for permission in result['Permissions']['Permission']:
                    if permission.get('Description') == description:
                        matched_rules.append(permission)
                        
            return matched_rules
        except Exception as e:
            logger.error(f"获取安全组规则失败: {e}")
            return []
            
    def remove_rules(self, security_group_id: str, rules: list):
        """
        删除指定的安全组规则
        
        Args:
            security_group_id: 安全组ID
            rules: 要删除的规则列表
        """
        for rule in rules:
            try:
                request = RevokeSecurityGroupRequest()
                request.set_SecurityGroupId(security_group_id)
                request.set_IpProtocol(rule['IpProtocol'])
                request.set_PortRange(rule['PortRange'])
                request.set_SourceCidrIp(rule['SourceCidrIp'])
                request.set_Policy(rule['Policy'])
                request.set_accept_format('json')
                
                self.client.do_action_with_exception(request)
                logger.info(f"成功删除阿里云安全组规则: {rule['SourceCidrIp']}")
            except Exception as e:
                logger.error(f"删除安全组规则失败: {e}")
                
    def add_rules(self, security_group_id: str, ips: set, port_range: str, description: str):
        """
        添加新的安全组规则
        
        Args:
            security_group_id: 安全组ID
            ips: IP地址集合
            port_range: 端口范围
            description: 规则描述
        """
        for ip in ips:
            try:
                request = AuthorizeSecurityGroupRequest()
                request.set_SecurityGroupId(security_group_id)
                request.set_IpProtocol('tcp')
                request.set_PortRange(port_range)
                request.set_SourceCidrIp(f"{ip}/32")
                request.set_Description(description)
                request.set_accept_format('json')
                
                self.client.do_action_with_exception(request)
                logger.info(f"成功添加阿里云安全组规则: {ip}")
            except Exception as e:
                logger.error(f"添加安全组规则失败: {e}")
                
    def sync_rules(self, security_group_id: str, ips: set, port_range: str = "1/65535", description: str = "Auto-synced rule"):
        """
        同步安全组规则
        
        Args:
            security_group_id: 安全组ID
            ips: 当前IP地址集合
            port_range: 端口范围
            description: 规则描述标识
        """
        # 获取现有规则
        existing_rules = self.get_existing_rules(security_group_id, description)
        logger.info(f"找到 {len(existing_rules)} 条现有规则需要删除")
        
        # 删除旧规则
        self.remove_rules(security_group_id, existing_rules)
        
        # 添加新规则
        logger.info(f"准备添加 {len(ips)} 条新规则")
        self.add_rules(security_group_id, ips, port_range, description)